Use dlopen to load NCCL EP - #3434
Conversation
Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>
Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>
1873e7b to
fdded06
Compare
Greptile SummaryThe PR moves NCCL EP from static linkage to on-demand runtime loading and packages its shared library and JIT headers with Transformer Engine. The revised native discovery now handles compatible fully versioned libraries, but Python availability detection still accepts incompatible-major files.
Confidence Score: 4/5The PR should not merge until Python availability detection rejects NCCL EP libraries that the native loader cannot load because of a major-version mismatch. An incompatible-major file under NCCL_EP_HOME currently passes the public availability checks but is excluded or rejected during native EP initialization. Files Needing Attention: transformer_engine/init.py and transformer_engine/common/ep/nccl_ep_provider.cpp Important Files Changed
Sequence DiagramsequenceDiagram
participant App
participant Python as Python availability check
participant Backend as EPBackend
participant Loader as NCCL EP loader
participant Library as libnccl_ep.so
App->>Python: is_nccl_ep_available()
Python->>Python: "Match NCCL_EP_HOME/lib/libnccl_ep.so*"
Python-->>App: Available
App->>Backend: Initialize EP
Backend->>Backend: Validate NCCL runtime
Backend->>Loader: initialize()
Loader->>Loader: Select matching NCCL_EP_MAJOR candidate
Loader->>Library: dlopen and validate version
Library-->>Backend: Symbols or load error
Reviews (3): Last reviewed commit: "switch provider to use `call_symbol`" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>
| shutil.rmtree(target_nccl_ep_dir) | ||
| shutil.copytree( | ||
| nccl_ep_dir, | ||
| target_nccl_ep_dir, | ||
| ) |
There was a problem hiding this comment.
Do we copy the whole nccl_ep dir instead of header files only?
There was a problem hiding this comment.
The whole built dir, but not the entire source repo. Here's what should be shipped:
transformer_engine/wheel_lib/
├── libtransformer_engine.so
├── libnccl_ep.so
└── nccl_ep/
└── include/
├── nccl_ep.h
└── nccl_ep/
├── config.h
├── common.hpp
├── ep_enums.h
└── device/
└── *.cuh
|
|
||
|
|
||
| def test_nccl_ep_library_not_found(monkeypatch): | ||
| monkeypatch.delenv("NCCL_EP_HOME", raising=False) |
There was a problem hiding this comment.
Does it mean NCCL_EP_HOME needs to be set at runtime?
There was a problem hiding this comment.
It does, but it's set automatically via nccl_ep_provider.cpp here. This allows users to continue using the TE shared lib directly without needing to explicitly set the var.
Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>
Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>
Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>
|
/te-ci |
2 participants
Description
NCCL EP is currently linked statically into the TE core shared library. This PR moves from statically linking to
dlopening the file at runtime, but only when NCCL EP features are needed.This solves two problems:
-z nowflag, which resolves all symbols at dynamic load time. This means that even if NCCL EP isn't used, it's dependencies symbols are still resolved. Because NCCL EP features require a newer version of NCCL than the rest of TE, this means that enabling BOLT optimization raises the minimum NCCL version. Bydlopening NCCL EP instead, resolution of these symbols controlled by TE, and can be limited to only be loaded when NCCL EP features are used.libandincludeare now shipped with TE as a part of this work.Type of change
Changes
Please list the changes introduced in this PR:
Checklist: